package com.example.scribe.backend.handlers;

import com.example.scribe.SyntaxNode;
import com.example.scribe.backend.ASTUtils;
import com.example.scribe.backend.BytecodeOps;
import com.example.scribe.backend.ExpressionEmitter;
import com.example.scribe.backend.LocalVariableAllocator;
import com.example.scribe.backend.StatementEmitter;
import com.example.scribe.backend.StatementNodeHandler;
import net.bytebuddy.implementation.bytecode.StackManipulation;

/**
 * Handler for if statements with conditional branching.
 */
public final class IfStatementHandler implements StatementNodeHandler {
    public IfStatementHandler() {
    }

    @Override
    public boolean canHandle(String nodeType) {
        return "IfStatement".equals(nodeType);
    }

    @Override
    public StackManipulation handle(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals, ExpressionEmitter exprEmitter, StatementEmitter statementEmitter) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        SyntaxNode condNode = ASTUtils.findFirst(node, "Condition");
        SyntaxNode thenNode = ASTUtils.findFirst(node, "Consequent");
        SyntaxNode elseNode = ASTUtils.findFirst(node, "Alternative");
        SyntaxNode condExpr = condNode != null ? ASTUtils.findFirst(condNode, "expression") : null;
        SyntaxNode thenBlock = thenNode != null ? ASTUtils.findFirst(thenNode, "block_stmt") : null;
        SyntaxNode elseBlock = elseNode != null ? ASTUtils.findFirst(elseNode, "block_stmt") : null;

        if (condExpr != null && thenBlock != null) {
            SyntaxNode thenRet = ASTUtils.extractReturnExpr(thenBlock);
            SyntaxNode elseRet = elseBlock != null ? ASTUtils.extractReturnExpr(elseBlock) : null;

            // Ensure condition expression is properly emitted (might be equality/comparison)
            // First try to find comparison/equality nodes directly
            StackManipulation cond = StackManipulation.Trivial.INSTANCE;
            SyntaxNode comparisonNode = ASTUtils.findFirst(condExpr, "comparison");
            if (comparisonNode == null) {
                comparisonNode = ASTUtils.findFirst(condExpr, "equality");
            }
            if (comparisonNode != null) {
                // Directly emit comparison node
                cond = exprEmitter.emit(comparisonNode, generatedClassName, locals);
            }
            if (cond == StackManipulation.Trivial.INSTANCE) {
                // Try emitting the condition expression itself
                cond = exprEmitter.emit(condExpr, generatedClassName, locals);
            }
            if (cond == StackManipulation.Trivial.INSTANCE) {
                // Fallback: try to find a simple value in the condition
                cond = BytecodeOps.pushInt(1); // default to true if can't evaluate
            }

            // Emit then branch return value - use similar logic to ReturnStatementHandler
            StackManipulation thenEmit;
            if (thenRet != null) {
                thenEmit = exprEmitter.emit(thenRet, generatedClassName, locals);
                if (thenEmit == StackManipulation.Trivial.INSTANCE) {
                    // Try to find identifier or number recursively
                    thenEmit = findValueInExpression(thenRet, generatedClassName, locals, exprEmitter);
                    if (thenEmit == StackManipulation.Trivial.INSTANCE) {
                        thenEmit = BytecodeOps.pushInt(0);
                    }
                }
            } else {
                thenEmit = BytecodeOps.pushInt(0);
            }

            // Emit else branch return value
            StackManipulation elseEmit;
            if (elseRet != null) {
                elseEmit = exprEmitter.emit(elseRet, generatedClassName, locals);
                if (elseEmit == StackManipulation.Trivial.INSTANCE) {
                    // Try to find identifier or number recursively
                    elseEmit = findValueInExpression(elseRet, generatedClassName, locals, exprEmitter);
                    if (elseEmit == StackManipulation.Trivial.INSTANCE) {
                        elseEmit = BytecodeOps.pushInt(0);
                    }
                }
            } else {
                elseEmit = BytecodeOps.pushInt(0);
            }

            return createBranch(cond, thenEmit, elseEmit, locals);
        }

        return StackManipulation.Trivial.INSTANCE;
    }

    /**
     * Create branching bytecode for if-then-else with proper stackmap frames.
     */
    private StackManipulation createBranch(StackManipulation cond, StackManipulation thenEmit, StackManipulation elseEmit, LocalVariableAllocator locals) {
        final StackManipulation thenEmitF = thenEmit;
        final StackManipulation elseEmitF = elseEmit;
        final int numLocals = locals.getMaxLocals(0);

        // Wrap condition to establish frame with locals before comparison
        StackManipulation condWithFrame = numLocals > 0 ? new StackManipulation.Compound(
            new StackManipulation.AbstractBase() {
                @Override
                public StackManipulation.Size apply(net.bytebuddy.jar.asm.MethodVisitor mv, net.bytebuddy.implementation.Implementation.Context ctx) {
                    // Establish frame with all local variables before condition evaluation
                    Object[] localTypes = new Object[numLocals];
                    java.util.Arrays.fill(localTypes, net.bytebuddy.jar.asm.Opcodes.INTEGER);
                    mv.visitFrame(net.bytebuddy.jar.asm.Opcodes.F_FULL, numLocals, localTypes, 0, null);
                    return new StackManipulation.Size(0, 0);
                }
            },
            cond
        ) : cond;

        return new StackManipulation.Compound(
            condWithFrame,
            new StackManipulation.AbstractBase() {
                @Override
                public StackManipulation.Size apply(net.bytebuddy.jar.asm.MethodVisitor mv, net.bytebuddy.implementation.Implementation.Context ctx) {
                    net.bytebuddy.jar.asm.Label thenL = new net.bytebuddy.jar.asm.Label();
                    net.bytebuddy.jar.asm.Label endL = new net.bytebuddy.jar.asm.Label();
                    mv.visitJumpInsn(net.bytebuddy.jar.asm.Opcodes.IFNE, thenL);
                    elseEmitF.apply(mv, ctx);
                    mv.visitInsn(net.bytebuddy.jar.asm.Opcodes.IRETURN);
                    mv.visitLabel(thenL);
                    // Frame at jump target (IFNE): condition consumed, stack empty
                    // Since we established locals frame before condition, use F_SAME to preserve them
                    mv.visitFrame(net.bytebuddy.jar.asm.Opcodes.F_SAME, 0, null, 0, null);
                    thenEmitF.apply(mv, ctx);
                    mv.visitInsn(net.bytebuddy.jar.asm.Opcodes.IRETURN);
                    mv.visitLabel(endL);
                    // endL is unreachable, no frame needed
                    return new StackManipulation.Size(0, 0);
                }
            }
        );
    }

    /**
     * Helper method to find value (identifier or number) in expression subtree.
     * Similar to findIdentifierOrNumber in ReturnStatementHandler.
     * This method prioritizes finding literal values over evaluating complex expressions.
     */
    private StackManipulation findValueInExpression(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals, ExpressionEmitter exprEmitter) {
        if (node == null) return StackManipulation.Trivial.INSTANCE;

        // Direct check for number first (highest priority)
        if (ASTUtils.isNumber(node)) {
            try {
                int v = Integer.parseInt(node.text());
                return BytecodeOps.pushInt(v);
            } catch (NumberFormatException e) {
                return StackManipulation.Trivial.INSTANCE;
            }
        }

        // Direct check for identifier
        if (ASTUtils.isIdentifier(node)) {
            String name = node.text();
            if (name != null && locals.hasName(name)) {
                return BytecodeOps.loadInt(locals.getIndexByName(name));
            }
        }

        // For expression nodes, recursively search children to find literals
        // This is important because expression nodes may wrap the actual literal values
        if ("expression".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if (!child.type().startsWith("Terminal:")) {
                    StackManipulation found = findValueInExpression(child, generatedClassName, locals, exprEmitter);
                    if (found != StackManipulation.Trivial.INSTANCE) {
                        return found;
                    }
                }
            }
        }

        // Check Call nodes FIRST - they may wrap primary nodes with numbers/identifiers
        // Structure: Call -> primary -> Terminal:NUMBER
        if ("Call".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if ("primary".equals(child.type())) {
                    for (SyntaxNode prim : child.children()) {
                        if (ASTUtils.isNumber(prim)) {
                            try {
                                int v = Integer.parseInt(prim.text());
                                return BytecodeOps.pushInt(v);
                            } catch (NumberFormatException e) {
                                // continue
                            }
                        }
                        if (ASTUtils.isIdentifier(prim)) {
                            String name = prim.text();
                            if (name != null && locals.hasName(name)) {
                                return BytecodeOps.loadInt(locals.getIndexByName(name));
                            }
                        }
                    }
                }
                // Recursively search all Call children
                StackManipulation found = findValueInExpression(child, generatedClassName, locals, exprEmitter);
                if (found != StackManipulation.Trivial.INSTANCE) {
                    return found;
                }
            }
        }

        // Check primary nodes for identifier or number
        if ("primary".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if (ASTUtils.isNumber(child)) {
                    try {
                        int v = Integer.parseInt(child.text());
                        return BytecodeOps.pushInt(v);
                    } catch (NumberFormatException e) {
                        // continue
                    }
                }
                if (ASTUtils.isIdentifier(child)) {
                    String name = child.text();
                    if (name != null && locals.hasName(name)) {
                        return BytecodeOps.loadInt(locals.getIndexByName(name));
                    }
                }
            }
            // Also recursively search primary children
            for (SyntaxNode child : node.children()) {
                StackManipulation found = findValueInExpression(child, generatedClassName, locals, exprEmitter);
                if (found != StackManipulation.Trivial.INSTANCE) {
                    return found;
                }
            }
        }

        // For structured nodes (shift, multiplication, unary, bitwise, etc.), recursively search children
        // These nodes may wrap the actual literal values
        // Important: bitwise_or, bitwise_xor, bitwise_and may have only one child (no operator)
        if ("shift".equals(node.type()) || "multiplication".equals(node.type()) || 
            "unary".equals(node.type()) || "addition".equals(node.type()) ||
            "subtraction".equals(node.type()) || "division".equals(node.type()) ||
            "bitwise_or".equals(node.type()) || "bitwise_xor".equals(node.type()) ||
            "bitwise_and".equals(node.type())) {
            // If only one child (no operator), recursively search it
            if (node.children().size() == 1) {
                return findValueInExpression(node.children().get(0), generatedClassName, locals, exprEmitter);
            } else {
                // Multiple children - skip operator terminals and recursively search
                for (SyntaxNode child : node.children()) {
                    if (!child.type().startsWith("Terminal:")) {
                        StackManipulation found = findValueInExpression(child, generatedClassName, locals, exprEmitter);
                        if (found != StackManipulation.Trivial.INSTANCE) {
                            return found;
                        }
                    }
                }
            }
        }

        // Recursively search all children (skip operator terminals)
        for (SyntaxNode child : node.children()) {
            if (!child.type().startsWith("Terminal:")) {
                StackManipulation found = findValueInExpression(child, generatedClassName, locals, exprEmitter);
                if (found != StackManipulation.Trivial.INSTANCE) {
                    return found;
                }
            }
        }

        return StackManipulation.Trivial.INSTANCE;
    }
}

